package com.kama.client.netty;

import common.trace.TraceContext;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.channel.ChannelOutboundHandlerAdapter;
import io.netty.channel.ChannelPromise;
import io.netty.util.AttributeKey;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;

/**
 * @author wxx
 * @version 1.0
 * @create 2025/2/20 16:56
 */
@Slf4j
// 继承自 ChannelOutboundHandlerAdapter 的自定义 ChannelHandler，用于处理 Netty 通道（Channel）的消息传输过程中的 Trace 上下文绑定
public class MDCChannelHandler extends ChannelOutboundHandlerAdapter {
    // 标识在 Netty 通道中存储 Trace 上下文的键
    public static final AttributeKey<Map<String, String>> TRACE_CONTEXT_KEY = AttributeKey.valueOf("TraceContext");
    // 将 Trace 上下文从通道属性提取并绑定到当前线程的 MDC 中，以便于后续的日志记录能够跟踪请求的追踪信息
    @Override
    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
        // 从当前通道 ctx.channel() 获取其属性 TRACE_CONTEXT_KEY 对应的值
        Map<String, String> traceContext =
                ctx.channel().attr(TRACE_CONTEXT_KEY).get();
        // 如果 traceContext 不为空，表示该通道已有 Trace 上下文信息
        if (traceContext != null) {
            // 设置到当前线程的TraceContext或MDC
            TraceContext.clone(traceContext);
            log.info("已绑定Trace上下文: {}", traceContext);
        } else {
            // 如果 traceContext 为空，则记录错误日志，表示没有找到 Trace 上下文
            log.error("Trace上下文未设置!");
        }

        // 继续传递请求
        super.write(ctx, msg, promise);
    }

}